import tkinter as tk
from tkinter import filedialog, messagebox, Checkbutton, IntVar
import pandas as pd
from Cryptodome.Cipher import AES
import base64
import re


class ExcelReader:
    def __init__(self, master):
        self.master = master
        self.file_path = tk.StringVar()
        self.df = pd.DataFrame()

        # 设置窗口标题和大小
        self.master.title("AES解密工具")
        self.master.geometry("600x500")
        # 900*750

        # 创建文件选择按钮
        self.btn_browse = tk.Button(
            self.master,
            text="选择文件",
            command=self.browse_file,
            width=15,
            height=2
        )
        self.btn_browse.pack(pady=20)

        # 显示选择的文件路径
        self.lbl_path = tk.Label(self.master)
        self.lbl_path.pack(pady=10)

        # 创建选择选项的容器
        self.options_frame = tk.Frame(self.master)
        self.options_frame.pack(pady=20)

        # 解密方式选择
        self.decrypt_choice = IntVar()
        self.rb_all = tk.Radiobutton(self.options_frame, text="解密所有内容", variable=self.decrypt_choice, value=1,
                                     command=self.hide_columns_options)
        self.rb_all.grid(row=0, column=0, padx=50)
        self.rb_select = tk.Radiobutton(self.options_frame, text="选择列进行解密", variable=self.decrypt_choice,
                                        value=2,
                                        command=self.show_columns_options)
        self.rb_select.grid(row=0, column=1, padx=50)

        # 创建文件字段标题
        self.lbl_file_fields = tk.Label(self.master, text="请勾选解密字段", font=("SimHei", 12, "bold"))

        # 存放字段的容器
        self.columns_frame = tk.Frame(self.master)

        # 为字段创建一个存储其状态的字典
        self.column_vars = {}
        self.decryption_choice = tk.IntVar()

        # 创建读取数据按钮
        self.btn_read_data = tk.Button(
            self.master,
            text="数据解密",
            command=self.read_data,
            width=15,
            height=2
        )
        self.btn_read_data.pack(side=tk.BOTTOM, pady=20)

    def hide_columns_options(self):
        for var in self.column_vars.values():
            var.set(0)

        self.lbl_file_fields.pack_forget()
        for widget in self.columns_frame.winfo_children():
            widget.destroy()
        self.columns_frame.pack_forget()

        self.master.geometry("600x500")

    def show_columns_options(self):
        self.lbl_file_fields.pack(pady=10)
        self.columns_frame.pack(pady=10)
        num_columns = len(self.df.columns)
        rows_needed = (num_columns + 3) // 4  # 四舍五入到最接近的整数

        # 根据字段数量调整窗口大小
        new_height = 500 + (rows_needed - 1) * 20  # 默认大小 + 额外的行
        self.master.geometry(f"600x{new_height}")

        for i, col in enumerate(self.df.columns):
            self.column_vars[col] = tk.IntVar()
            Checkbutton(self.columns_frame, text=col, variable=self.column_vars[col]).grid(row=i // 4, column=i % 4,
                                                                                           sticky='w', padx=20)

    def browse_file(self):
        file_path = filedialog.askopenfilename(
            initialdir="/",
            title="选择文件",
            filetypes=(("Excel files", "*.xls;*.xlsx"), ("all files", "*.*"))
        )
        if file_path:
            self.file_path.set(file_path)
            self.lbl_path.config(text=self.file_path.get())
            self.df = pd.read_excel(self.file_path.get(), dtype=str)

    def read_data(self):
        # 获取选中的字段
        selected_columns = [col for col, var in self.column_vars.items() if var.get()]

        # 根据选择执行解密操作
        if self.decrypt_choice.get() == 1:  # 解密所有内容
            self.decrypt_all()
        elif self.decrypt_choice.get() == 2:  # 解密部分内容
            self.decrypt_selected(selected_columns)
        else:
            messagebox.showwarning("提示", "请先选择解密方式")
            return

        self.prompt_save_data()

    def decrypt_all(self):
        pattern = r'^(?=.*[a-z])(?=.*[A-Z])(?=.*\d)[\w+/=]*$'

        original_columns = [col for col in self.df.columns if "_解密后" not in col]

        for col in original_columns:
            self.df[col] = self.df[col].apply(
                lambda x: _decrypt(x) if pd.notna(x) and re.match(pattern, str(x)) else x
            )

        for col in list(self.df.columns):
            if "_解密后" in col:
                self.df.drop(columns=col, inplace=True)

    def decrypt_selected(self, columns):
        pattern = r'^(?=.*[a-z])(?=.*[A-Z])(?=.*\d)[\w+/=]*$'
        for col in columns:
            decrypted_col = self.df[col].apply(
                lambda x: _decrypt(x) if pd.notna(x) and re.match(pattern, str(x)) else x)
            self.df[f"{col}_解密后"] = decrypted_col

    def prompt_save_data(self):
        answer = messagebox.askquestion("成功", "数据解密成功！是否保存当前文件？")
        if answer == 'yes':
            self.save_data()

    def save_data(self):
        try:
            save_file_path = filedialog.asksaveasfilename(
                defaultextension=".xlsx",
                filetypes=(("Excel files", "*.xlsx"), ("all files", "*.*"))
            )
            if not save_file_path:
                return
            self.df.to_excel(save_file_path, index=False)
            messagebox.showinfo("成功", "数据保存成功！")
        except Exception as e:
            messagebox.showerror("错误", "数据保存错误：" + str(e))


def aes_ecb_decrypt(ciphertext, key):
    cipher = AES.new(key, AES.MODE_ECB)
    plaintext = cipher.decrypt(ciphertext)
    return plaintext


def _decrypt(x):
    key = b""  # AES加密密钥
    try:
        de_result = aes_ecb_decrypt(base64.b64decode(x.encode()), key).decode()
        de_result = de_result.replace('', '').replace('\x00', '')
    except Exception as e:
        de_result = None
    return de_result


if __name__ == '__main__':
    root = tk.Tk()
    ExcelReader(root)
    root.mainloop()



